#!/usr/bin/env python
# encoding: utf-8
from __future__ import division
from PIL import Image
import os
import os.path
from xml.etree.ElementTree import parse, Element
import sys
import random
import cPickle
reload(sys)
sys.setdefaultencoding('utf8')
lst = []
cls_name = 'toothbrush'

def pair_box_jpg(res_xml, jpg):
    """corresponding candidate boxes and pictures,which save in a list
    like [xmin,ymin,xmax,ymax,path]
    """
    doc = parse(res_xml)
    root = doc.getroot()
    for object in root.iter('object'):
        name_context = object.find('name').text
        if name_context != cls_name:
            continue
        bbox = object.find('bndbox')
        xmin = int(bbox.find('xmin').text)
        ymin = int(bbox.find('ymin').text)
        xmax = int(bbox.find('xmax').text)
        ymax = int(bbox.find('ymax').text)
        area = (xmax - xmin) * (ymax - ymin)
        # you can specify your own conditions to filter boxes
        if area > 10000 and area < 50000:
            lst.append([xmin, ymin, xmax, ymax, jpg])


def get_box_region(base_img):
    width, height = base_img.size
    while(1):
        # select one box for candidate box
        index = random.randint(0, len(lst) - 1)
        [xmin, ymin, xmax, ymax, path] = lst[index]
        box_width = xmax - xmin
        box_height = ymax - ymin

        # Whether base image can contain a selected box
        if width > box_width and height > box_height:
            break

    tmp_img = Image.open(path)
    box = tmp_img.crop((xmin, ymin, xmax, ymax))

    # Calculate a region to place the selected box,region and selected box is
    # the same size
    region_xmin = random.randint(0, width - box_width)
    region_ymin = random.randint(0, height - box_height)
    region_xmax = region_xmin + box_width
    region_ymax = region_ymin + box_height

    region = (region_xmin, region_ymin, region_xmax, region_ymax)

    return box, region


def cross_line(rect1, rect2):
    """whether two rectangle are croessed
    """
    xmin1, ymin1, xmax1, ymax1 = rect1
    xmin2, ymin2, xmax2, ymax2 = rect2
    flag1 = abs((xmin1 + xmax1) / 2 - (xmin2 + xmax2) /
                2) < ((xmax1 + xmax2 - xmin1 - xmin2) / 2)
    flag2 = abs((ymin1 + ymax1) / 2 - (ymin2 + ymax2) /
                2) < ((ymax1 + ymax2 - ymin1 - ymin2) / 2)
    if flag1 and flag2:
        return True
    else:
        return False


def main():
    with open('./th.txt', 'r') as f:
        lines = f.readlines()

    for line in lines:
        xml = './Annotations/' + line.strip() + '.xml'
        jpg = '/data2/ad_detect/tmp/Deformable-ConvNets2/data/VOCdevkit/VOC2007/JPEGImages/' + line.strip() + \
            '.jpg'
        pair_box_jpg(xml, jpg)
    # base image put here
    rootdir = './output'
    annotations = {}
    for parent, dirnames, filenames in os.walk(rootdir):
        for filename in filenames:
            base_img = Image.open(os.path.join(parent, filename))
            # paste two box on image, ensure two region not crossed
            box, region = get_box_region(base_img)
            while(1):
                box1, region1 = get_box_region(base_img)
                if cross_line(region, region1):
                    continue
                else:
                    break
            base_img.paste(box, region)
            base_img.paste(box1, region1)
            annotations[filename] = [region, region1]
            # pase image after pasting
            base_img.save('./tmp/' + filename)
    # record image and two box by a dictionary, and which serialized by a pkl file
    cPickle.dump(annotations, open('anno.pkl', 'wb'))


if __name__ == '__main__':
    main()
